using System;
using System.Collections.Generic;

namespace LaplaceMeshDeformer
{
    /// <summary>
    /// 稀疏矩阵类
    /// 使用行压缩存储(CSR)格式存储稀疏矩阵
    /// </summary>
    public class SparseMatrix
    {
        private int rows;           // 行数
        private int columns;        // 列数
        private List<double> values;    // 非零元素值
        private List<int> columnIndices; // 非零元素的列索引
        private List<int> rowPointers;   // 每行起始位置的指针

        /// <summary>
        /// 获取矩阵的行数
        /// </summary>
        public int Rows => rows;

        /// <summary>
        /// 获取矩阵的列数
        /// </summary>
        public int Columns => columns;

        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="rows">行数</param>
        /// <param name="columns">列数</param>
        public SparseMatrix(int rows, int columns)
        {
            if (rows <= 0 || columns <= 0)
            {
                throw new ArgumentException("矩阵维度必须为正数");
            }

            this.rows = rows;
            this.columns = columns;
            this.values = new List<double>();
            this.columnIndices = new List<int>();
            this.rowPointers = new List<int>();

            // 初始化行指针
            for (int i = 0; i <= rows; i++)
            {
                rowPointers.Add(0);
            }
        }

        /// <summary>
        /// 设置矩阵元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="column">列索引</param>
        /// <param name="value">元素值</param>
        public void SetElement(int row, int column, double value)
        {
            if (row < 0 || row >= rows || column < 0 || column >= columns)
            {
                throw new ArgumentOutOfRangeException("索引超出矩阵范围");
            }

            // 查找元素在稀疏存储中的位置
            int startPos = rowPointers[row];
            int endPos = rowPointers[row + 1];
            int insertPos = -1;

            for (int i = startPos; i < endPos; i++)
            {
                if (columnIndices[i] == column)
                {
                    // 找到元素位置，直接修改值
                    values[i] = value;
                    return;
                }
                else if (columnIndices[i] > column && insertPos == -1)
                {
                    // 找到插入位置
                    insertPos = i;
                }
            }

            // 如果值接近于0，且不存在该元素，则不添加
            if (Math.Abs(value) < 1e-10)
            {
                return;
            }

            // 如果没有找到元素，需要插入新元素
            if (insertPos == -1)
            {
                // 插入到该行的末尾
                insertPos = endPos;
            }

            // 插入新元素
            values.Insert(insertPos, value);
            columnIndices.Insert(insertPos, column);

            // 更新后续行的行指针
            for (int i = row + 1; i <= rows; i++)
            {
                rowPointers[i]++;
            }
        }

        /// <summary>
        /// 获取矩阵元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="column">列索引</param>
        /// <returns>元素值</returns>
        public double GetElement(int row, int column)
        {
            if (row < 0 || row >= rows || column < 0 || column >= columns)
            {
                throw new ArgumentOutOfRangeException("索引超出矩阵范围");
            }

            // 查找元素在稀疏存储中的位置
            int startPos = rowPointers[row];
            int endPos = rowPointers[row + 1];

            for (int i = startPos; i < endPos; i++)
            {
                if (columnIndices[i] == column)
                {
                    return values[i];
                }
            }

            // 没找到元素，返回0
            return 0.0;
        }

        /// <summary>
        /// 矩阵向量乘法：y = Ax
        /// </summary>
        /// <param name="x">输入向量</param>
        /// <param name="y">输出向量</param>
        public void MultiplyVector(double[] x, double[] y)
        {
            if (x.Length != columns || y.Length != rows)
            {
                throw new ArgumentException("向量维度与矩阵不匹配");
            }

            // 初始化输出向量为0
            for (int i = 0; i < rows; i++)
            {
                y[i] = 0.0;
            }

            // 执行稀疏矩阵-向量乘法
            for (int i = 0; i < rows; i++)
            {
                int startPos = rowPointers[i];
                int endPos = rowPointers[i + 1];

                for (int j = startPos; j < endPos; j++)
                {
                    int col = columnIndices[j];
                    double val = values[j];
                    y[i] += val * x[col];
                }
            }
        }

        /// <summary>
        /// 清除矩阵中的所有元素
        /// </summary>
        public void Clear()
        {
            values.Clear();
            columnIndices.Clear();

            // 重置行指针
            rowPointers.Clear();
            for (int i = 0; i <= rows; i++)
            {
                rowPointers.Add(0);
            }
        }

        /// <summary>
        /// 获取矩阵的非零元素数量
        /// </summary>
        /// <returns>非零元素数量</returns>
        public int NonZeroCount()
        {
            return values.Count;
        }

        /// <summary>
        /// 转置矩阵
        /// </summary>
        /// <returns>转置后的矩阵</returns>
        public SparseMatrix Transpose()
        {
            SparseMatrix result = new SparseMatrix(columns, rows);

            // 计算每一列（转置后的每一行）的非零元素数量
            int[] colCounts = new int[columns];
            for (int i = 0; i < columnIndices.Count; i++)
            {
                colCounts[columnIndices[i]]++;
            }

            // 设置转置矩阵的行指针
            result.rowPointers[0] = 0;
            for (int i = 0; i < columns; i++)
            {
                result.rowPointers[i + 1] = result.rowPointers[i] + colCounts[i];
            }

            // 分配空间
            int nnz = values.Count;
            result.values = new List<double>(new double[nnz]);
            result.columnIndices = new List<int>(new int[nnz]);

            // 复制元素 (CSR -> CSR 转置)
            int[] nextPos = new int[columns];
            for (int i = 0; i < columns; i++)
            {
                nextPos[i] = result.rowPointers[i];
            }

            for (int i = 0; i < rows; i++)
            {
                for (int j = rowPointers[i]; j < rowPointers[i + 1]; j++)
                {
                    int col = columnIndices[j];
                    int pos = nextPos[col]++;
                    result.columnIndices[pos] = i;
                    result.values[pos] = values[j];
                }
            }

            return result;
        }

        /// <summary>
        /// 创建单位矩阵
        /// </summary>
        /// <param name="size">矩阵大小</param>
        /// <returns>单位矩阵</returns>
        public static SparseMatrix Identity(int size)
        {
            SparseMatrix result = new SparseMatrix(size, size);

            for (int i = 0; i < size; i++)
            {
                result.SetElement(i, i, 1.0);
            }

            return result;
        }

        /// <summary>
        /// 返回矩阵信息的字符串表示
        /// </summary>
        /// <returns>矩阵信息字符串</returns>
        public override string ToString()
        {
            return $"SparseMatrix({rows}x{columns}, 非零元素: {NonZeroCount()})";
        }
    }
} 